from annoy import AnnoyIndex
import numpy as np
import os

# 1. 模拟图书馆数据库
np.random.seed(42)
num_books = 100000  # 模拟10万本图书
vector_dimension = 128  # 每本书的特征向量为128维
book_vectors = np.random.rand(num_books, vector_dimension)  # 生成随机特征向量

# 假设每本书的ID为0~99999
book_ids = np.arange(num_books)

# 模拟用户查询的特征向量
query_vector = np.random.rand(vector_dimension)

# 2. 构建Annoy索引
num_trees = 50  # 随机投影树数量
index_path = "library_annoy_index.ann"  # 持久化索引文件路径

# 创建Annoy索引
annoy_index = AnnoyIndex(vector_dimension, metric='euclidean')

# 添加图书向量到索引中
for i in range(num_books):
    annoy_index.add_item(i, book_vectors[i])

# 构建索引
print("构建索引中，请稍候...")
annoy_index.build(num_trees)

# 持久化索引到磁盘
annoy_index.save(index_path)
print("索引已保存到磁盘。")

# 3. 查询图书
loaded_index = AnnoyIndex(vector_dimension, metric='euclidean')
loaded_index.load(index_path)  # 从磁盘加载索引

# 查询与用户输入相似的5本图书
top_k = 5
search_k = 2000  # 控制搜索树数量，平衡速度与精度
results = loaded_index.get_nns_by_vector(query_vector, top_k, search_k=search_k, include_distances=True)

# 4. 输出查询结果
print("\n用户查询特征向量:")
print(query_vector)
print("\n检索到的相似图书:")
for idx, (book_id, distance) in enumerate(zip(results[0], results[1])):
    print(f"结果{idx + 1}: 图书ID {book_id}, 距离 {distance:.4f}, 特征向量 {book_vectors[book_id]}")

# 5. 性能分析
index_size = os.path.getsize(index_path) / (1024 * 1024)  # 索引大小(MB)
print("\n性能分析:")
print(f"图书总数: {num_books}")
print(f"索引文件大小: {index_size:.2f} MB")
print(f"随机投影树数量: {num_trees}")
print(f"查询使用的树数量: {search_k}")

# 删除持久化索引文件（清理环境）
os.remove(index_path)